use crate::common::error_enum::ErrorEnum;
use crate::packet_hdr::{
    ApplicationType, EtherType, IpAddr, IpHeader, MacAddr, NetworkPacket, TransportType,
};
use crate::pcap_parse::pcap_hdr::*;
use std::collections::HashMap;

use crate::common::error_enum::ErrorEnum::NoPcapIndexError;
use log::{error, info};
use std::fs::File;
use std::io::SeekFrom::Start;
use std::io::{BufReader, Read, Seek};

trait SortPacket {
    fn sort_by_ethernet_type(&mut self, range: &Vec<usize>) -> HashMap<EtherType, Vec<usize>>;
    fn sort_by_transport_type(&mut self, range: &Vec<usize>) -> HashMap<TransportType, Vec<usize>>;
    fn sort_by_application_type(
        &mut self,
        range: &Vec<usize>,
    ) -> HashMap<ApplicationType, Vec<usize>>;
    fn sort_by_ipaddr(&mut self, range: &Vec<usize>) -> (HashMap<IpAddr, Vec<usize>>);
    fn sort_by_mac(&mut self, range: &Vec<usize>) -> HashMap<MacAddr, Vec<usize>>;
    fn sort_by_src_mac(&mut self, range: &Vec<usize>) -> HashMap<MacAddr, Vec<usize>>;
    fn sort_by_dst_mac(&mut self, range: &Vec<usize>) -> HashMap<MacAddr, Vec<usize>>;
    fn sort_by_src_ip(&mut self, range: &Vec<usize>) -> (HashMap<IpAddr, Vec<usize>>);
    fn sort_by_dst_ip(&mut self, range: &Vec<usize>) -> (HashMap<IpAddr, Vec<usize>>);
    fn sort_by_port(&mut self, range: &Vec<usize>) -> HashMap<isize, Vec<usize>>;
    fn sort_by_src_port(&mut self, range: &Vec<usize>) -> HashMap<isize, Vec<usize>>;
    fn sort_by_dst_port(&mut self, range: &Vec<usize>) -> HashMap<isize, Vec<usize>>;
}

pub struct PcapFile {
    pub pcap_file_buf: BufReader<File>,
    pub offset: usize, //相对于pcap_hdr的偏移
    pub pcap_hdr: PcapHdr,
    pub endianness: Endianness,
    pub file_format: FileFormat,
    pcap_index_arr: Option<Vec<PcapIndex>>,
    pcap_num: usize,
}

impl SortPacket for PcapFile {
    fn sort_by_ethernet_type(&mut self, range: &Vec<usize>) -> HashMap<EtherType, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let ether_type = network_packet.as_ref().unwrap().get_ether_type();
            ret.entry(ether_type).or_insert(Vec::new()).push(*index);
        }
        ret
    }
    fn sort_by_transport_type(&mut self, range: &Vec<usize>) -> HashMap<TransportType, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let transport_type = network_packet.as_ref().unwrap().get_transport_type();
            ret.entry(transport_type).or_insert(Vec::new()).push(*index);
        }
        ret
    }
    fn sort_by_application_type(
        &mut self,
        range: &Vec<usize>,
    ) -> HashMap<ApplicationType, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let application_type = network_packet.as_ref().unwrap().get_application_type();
            ret.entry(application_type)
                .or_insert(Vec::new())
                .push(*index);
        }
        ret
    }

    fn sort_by_ipaddr(&mut self, range: &Vec<usize>) -> (HashMap<IpAddr, Vec<usize>>) {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }

            let src_addr = network_packet.as_ref().unwrap().get_ip_src_addr();
            let dst_addr = network_packet.as_ref().unwrap().get_ip_dst_addr();
            let flag = src_addr == dst_addr;
            ret.entry(src_addr).or_insert(Vec::new()).push(*index);
            if flag {
                continue;
            }
            ret.entry(dst_addr).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_mac(&mut self, range: &Vec<usize>) -> HashMap<MacAddr, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let dst_mac = network_packet.as_ref().unwrap().get_dst_mac();
            let src_mac = network_packet.as_ref().unwrap().get_src_mac();
            ret.entry(dst_mac).or_insert(Vec::new()).push(*index);
            if dst_mac == src_mac {
                continue;
            }
            ret.entry(src_mac).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_src_mac(&mut self, range: &Vec<usize>) -> HashMap<MacAddr, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }

            let src_mac = network_packet.as_ref().unwrap().get_src_mac();

            ret.entry(src_mac).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_dst_mac(&mut self, range: &Vec<usize>) -> HashMap<MacAddr, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let dst_mac = network_packet.as_ref().unwrap().get_dst_mac();

            ret.entry(dst_mac).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_src_ip(&mut self, range: &Vec<usize>) -> (HashMap<IpAddr, Vec<usize>>) {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let src_addr = network_packet.as_ref().unwrap().get_ip_src_addr();
            ret.entry(src_addr).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_dst_ip(&mut self, range: &Vec<usize>) -> (HashMap<IpAddr, Vec<usize>>) {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }
            let dst_addr = network_packet.as_ref().unwrap().get_ip_dst_addr();
            ret.entry(dst_addr).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_port(&mut self, range: &Vec<usize>) -> HashMap<isize, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }

            let src_port = network_packet.as_ref().unwrap().get_src_port();
            let dst_port = network_packet.as_ref().unwrap().get_dst_port();

            ret.entry(src_port).or_insert(Vec::new()).push(*index);
            if src_port == dst_port {
                continue;
            }
            ret.entry(dst_port).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_src_port(&mut self, range: &Vec<usize>) -> HashMap<isize, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }

            let src_port = network_packet.as_ref().unwrap().get_src_port();

            ret.entry(src_port).or_insert(Vec::new()).push(*index);
        }
        ret
    }

    fn sort_by_dst_port(&mut self, range: &Vec<usize>) -> HashMap<isize, Vec<usize>> {
        let mut ret = HashMap::new();
        if let None = self.pcap_index_arr {
            return ret;
        }
        for index in range {
            if *index >= self.pcap_num {
                continue;
            }
            let network_packet = NetworkPacket::new_from_pcap(*index, self);
            if let Err(_) = network_packet {
                continue;
            }

            let dst_port = network_packet
                .as_ref()
                .unwrap()
                .transport_header
                .get_dst_port();

            ret.entry(dst_port).or_insert(Vec::new()).push(*index);
        }
        ret
    }
}

impl PcapFile {
    pub fn new(file_path: &str) -> Result<Self, ErrorEnum> {
        let file = File::open(file_path)?;
        let mut pcap_file_buf = BufReader::new(file);
        let pcap_hdr = PcapHdr::new(&mut pcap_file_buf)?;
        let (endianness, file_format) = PcapFile::get_pcap_format(&mut pcap_file_buf)?;
        Ok(Self {
            pcap_file_buf,
            offset: PCAP_HDR_SIZE,
            pcap_hdr,
            endianness,
            file_format,
            pcap_index_arr: None,
            pcap_num: 0,
        })
    }
    pub fn build_index_from_pcap(&mut self) -> Result<(), ErrorEnum> {
        let mut arr = Vec::new();
        let mut buf = vec![0; PCAP_PRERECORD_HDR_SIZE];
        let mut offset = self.offset;
        self.pcap_file_buf.seek(Start(offset as u64))?;
        while PCAP_PRERECORD_HDR_SIZE == self.pcap_file_buf.read(&mut buf)? {
            let pcap_record_hdr =
                PcapRecordHdr::new(&mut buf, (self.endianness, self.file_format))?;
            let pcap_index = PcapIndex {
                index: offset,
                pcap_record_hdr,
                data: None,
            };
            offset += pcap_index.pcap_record_hdr.incl_len as usize + PCAP_PRERECORD_HDR_SIZE;

            self.pcap_file_buf.seek(Start(offset as u64))?;
            arr.push(pcap_index);
        }
        self.pcap_num = arr.len();
        self.pcap_index_arr = Some(arr);
        Ok(())
    }

    pub fn get_single_pcap_mut(&mut self, pcap_index: usize) -> Result<&mut PcapIndex, ErrorEnum> {
        if let None = self.pcap_index_arr {
            return Err(NoPcapIndexError(pcap_index));
        }
        match self.pcap_index_arr.as_mut() {
            None => Err(NoPcapIndexError(pcap_index)),
            Some(ret) => {
                if ret.len() <= pcap_index {
                    return Err(NoPcapIndexError(pcap_index));
                }
                Ok(&mut ret[pcap_index])
            }
        }
    }

    pub fn get_single_pcap(&self, pcap_index: usize) -> Result<&PcapIndex, ErrorEnum> {
        if let None = self.pcap_index_arr {
            return Err(NoPcapIndexError(pcap_index));
        }
        match self.pcap_index_arr.as_ref() {
            None => Err(NoPcapIndexError(pcap_index)),
            Some(ret) => {
                if ret.len() <= pcap_index {
                    return Err(NoPcapIndexError(pcap_index));
                }
                Ok(&ret[pcap_index])
            }
        }
    }

    pub fn cache(
        &mut self,
        pcap_index: usize,
        network_packet: NetworkPacket,
    ) -> &Option<Box<NetworkPacket>> {
        match self.get_single_pcap_mut(pcap_index) {
            Ok(pcap_index) => {
                pcap_index.data = Some(Box::new(network_packet));
                &pcap_index.data
            }
            Err(e) => {
                info!("cache network packet error : {:?}", e);
                &None
            }
        }
    }

    pub fn get_pcap_index_arr(&self) -> &Option<Vec<PcapIndex>> {
        &self.pcap_index_arr
    }

    pub fn get_single_pcap_bytes(&mut self, pcap_index: usize) -> Vec<u8> {
        || -> Result<Vec<u8>, ErrorEnum> {
            let pcap_index = self.get_single_pcap_mut(pcap_index)?;
            let index = pcap_index.index;
            let end = pcap_index.pcap_record_hdr.incl_len as usize + PCAP_PRERECORD_HDR_SIZE;
            self.pcap_file_buf.seek(Start(index as u64))?;
            let mut buf = vec![0; end];
            self.pcap_file_buf.read_exact(&mut buf)?;
            Ok(buf)
        }()
        .unwrap_or_else(|e| {
            error!("get bytes err: {}", e);
            Vec::new()
        })
    }
    pub fn get_index_num(&self) -> usize {
        self.pcap_num
    }

    pub fn get_network_packet(&self, pcap_index: usize) -> Option<Box<NetworkPacket>> {
        self.get_single_pcap(pcap_index)
            .map(|pcap_index| pcap_index.data.clone())
            .unwrap_or_else(|_| None)
    }

    pub fn get_pcap_format(
        pcap_file_buf: &mut BufReader<File>,
    ) -> Result<(Endianness, FileFormat), ErrorEnum> {
        pcap_file_buf.seek(Start(0))?;
        let mut buf = vec![0; 4];
        let bytes_read_num = pcap_file_buf.read(&mut buf)?;
        if bytes_read_num != 4 {
            return Err(ErrorEnum::FileError);
        }
        let magic_number = &buf[0..4];
        let magic_number = format!(
            "{:02X}{:02X}{:02X}{:02X}",
            magic_number[0], magic_number[1], magic_number[2], magic_number[3]
        );
        match &magic_number as &str {
            // 标准格式，大端序
            "A1B2C3D4" => Ok((Endianness::BigEndian, FileFormat::Standard)),
            // 标准格式，小端序
            "D4C3B2A1" => Ok((Endianness::LittleEndian, FileFormat::Standard)),
            // NSG 格式，大端序
            "A1B23C4D" => Ok((Endianness::BigEndian, FileFormat::Nsg)),
            // NSG 格式，小端序
            "4D3CB2A1" => Ok((Endianness::LittleEndian, FileFormat::Nsg)),
            // 无法识别的 magic number
            _ => Err(ErrorEnum::InvalidMagicNumber(magic_number)),
        }
    }
}

#[derive(Debug)]
pub struct PcapIndex {
    pub index: usize,
    pub pcap_record_hdr: PcapRecordHdr,
    pub data: Option<Box<NetworkPacket>>, //box空指针优化避免没有数据的时候占用大量空间
}

pub fn parse_u16(data: [u8; 2], endianness: Endianness) -> u16 {
    match endianness {
        Endianness::BigEndian => u16::from_be_bytes(data),
        Endianness::LittleEndian => u16::from_le_bytes(data),
    }
}

pub fn parse_u32(data: [u8; 4], endianness: Endianness) -> u32 {
    match endianness {
        Endianness::BigEndian => u32::from_be_bytes(data),
        Endianness::LittleEndian => u32::from_le_bytes(data),
    }
}

pub fn parse_i32(data: [u8; 4], endianness: Endianness) -> i32 {
    match endianness {
        Endianness::BigEndian => i32::from_be_bytes(data),
        Endianness::LittleEndian => i32::from_le_bytes(data),
    }
}
